# 导入必要的库
from transformers import pipeline, AutoTokenizer, AutoModelForMaskedLM,AutoModel
import torch
# 加载模型和tokenizer
tokenizers = AutoTokenizer.from_pretrained("bert-base-chinese")
model = AutoModelForMaskedLM.from_pretrained("bert-base-chinese")

# 创建一个pipeline对象，指定任务为fill-mask
fill_mask = pipeline("fill-mask", model=model, tokenizer=tokenizers)


# 定义一个函数，用于生成聊天机器人的回复
def chatbot_reply(user_input):
    while(True):
        text = user_input + '[MASK]'
        tokens = tokenizers(text, return_tensors="pt")
        token_logits = model(**tokens).logits
        # 找到MASK的位置然后获取其预测值
        mask_token_index = torch.where(tokens["input_ids"] == tokenizers.mask_token_id)[1]
        mask_token_logits = token_logits[0, mask_token_index, :]
        # 对MASK所在位置找到他的TOP5预测结果
        top_1_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()
        user_input = text.replace(tokenizers.mask_token, tokenizers.decode([top_1_tokens[0]]))
        if len(user_input)>100:
            return user_input


# 测试聊天机器人
user_input = "你好，我叫[MASK]龙"
reply = chatbot_reply(user_input)
print(reply)  # 输出：你好，我叫小明，你叫什么？我叫小红。
